"""数据库迁移脚本

为BrainWave表添加时序相关字段
支持数据包时间戳、时长、序列号等时序信息
"""

import sys
from datetime import datetime
from loguru import logger
from peewee import *
from playhouse.migrate import *

from app.config import settings
from app.core.database import database, init_database


def migrate_brain_wave_table():
    """迁移BrainWave表，添加时序字段"""
    try:
        logger.info("开始迁移BrainWave表...")
        
        # 连接数据库
        if not init_database():
            logger.error("数据库连接失败")
            return False
        
        # 检查表是否存在
        table_exists = False
        try:
            result = database.execute_sql("SHOW TABLES LIKE 'brain_wave'").fetchone()
            table_exists = result is not None
        except Exception as e:
            logger.error(f"检查表存在性失败: {str(e)}")
            return False
        
        if not table_exists:
            logger.info("BrainWave表不存在，请先运行主程序创建基础表结构")
            return False
        
        # 检查字段是否已存在
        columns = database.get_columns('brain_wave')
        existing_columns = [col.name for col in columns]
        
        new_fields = [
            'data_timestamp',
            'duration_seconds', 
            'sequence_number',
            'client_send_time',
            'updated_at'
        ]
        
        fields_to_add = [field for field in new_fields if field not in existing_columns]
        
        if not fields_to_add:
            logger.info("所有时序字段已存在，无需迁移")
            return True
        
        logger.info(f"需要添加的字段: {fields_to_add}")
        
        # 添加新字段
        migrator = MySQLMigrator(database)
        
        migrations = []
        if 'data_timestamp' in fields_to_add:
            migrations.append(
                migrator.add_column('brain_wave', 'data_timestamp', FloatField(null=True))
            )
            
        if 'duration_seconds' in fields_to_add:
            migrations.append(
                migrator.add_column('brain_wave', 'duration_seconds', FloatField(null=True))
            )
            
        if 'sequence_number' in fields_to_add:
            migrations.append(
                migrator.add_column('brain_wave', 'sequence_number', IntegerField(null=True))
            )
            
        if 'client_send_time' in fields_to_add:
            migrations.append(
                migrator.add_column('brain_wave', 'client_send_time', FloatField(null=True))
            )
            
        if 'updated_at' in fields_to_add:
            migrations.append(
                migrator.add_column('brain_wave', 'updated_at', DateTimeField(default=datetime.now))
            )
        
        # 执行迁移
        if migrations:
            database.execute_sql('SET FOREIGN_KEY_CHECKS = 0;')  # 临时禁用外键检查
            
            for migration in migrations:
                migrate(migration)
                
            database.execute_sql('SET FOREIGN_KEY_CHECKS = 1;')  # 重新启用外键检查
            
            logger.info(f"成功添加 {len(migrations)} 个字段到BrainWave表")
        
        # 验证迁移结果
        columns_after = database.get_columns('brain_wave')
        columns_after_names = [col.name for col in columns_after]
        
        for field in new_fields:
            if field in columns_after_names:
                logger.info(f"✓ 字段 {field} 添加成功")
            else:
                logger.error(f"✗ 字段 {field} 添加失败")
                return False
        
        logger.info("BrainWave表迁移完成")
        return True
        
    except Exception as e:
        logger.error(f"迁移BrainWave表失败: {str(e)}")
        return False


def rollback_migration():
    """回滚迁移，删除时序字段"""
    try:
        logger.info("开始回滚BrainWave表迁移...")
        
        if not init_database():
            logger.error("数据库连接失败")
            return False
        
        # 检查表是否存在
        table_exists = False
        try:
            result = database.execute_sql("SHOW TABLES LIKE 'brain_wave'").fetchone()
            table_exists = result is not None
        except Exception as e:
            logger.error(f"检查表存在性失败: {str(e)}")
            return False
        
        if not table_exists:
            logger.info("BrainWave表不存在，无需回滚")
            return True
        
        migrator = MySQLMigrator(database)
        
        fields_to_remove = [
            'data_timestamp',
            'duration_seconds',
            'sequence_number', 
            'client_send_time'
        ]
        
        # 检查字段是否存在
        columns = database.get_columns('brain_wave')
        existing_columns = [col.name for col in columns]
        
        fields_to_remove = [field for field in fields_to_remove if field in existing_columns]
        
        if not fields_to_remove:
            logger.info("没有需要删除的时序字段")
            return True
        
        logger.info(f"将删除字段: {fields_to_remove}")
        
        # 删除字段
        database.execute_sql('SET FOREIGN_KEY_CHECKS = 0;')
        
        for field in fields_to_remove:
            try:
                migrate(migrator.drop_column('brain_wave', field))
                logger.info(f"✓ 字段 {field} 删除成功")
            except Exception as e:
                logger.error(f"✗ 删除字段 {field} 失败: {str(e)}")
        
        database.execute_sql('SET FOREIGN_KEY_CHECKS = 1;')
        
        logger.info("BrainWave表回滚完成")
        return True
        
    except Exception as e:
        logger.error(f"回滚迁移失败: {str(e)}")
        return False


def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("用法: python migrate_database.py [migrate|rollback]")
        sys.exit(1)
    
    action = sys.argv[1].lower()
    
    if action == "migrate":
        success = migrate_brain_wave_table()
        if success:
            logger.info("数据库迁移成功")
            sys.exit(0)
        else:
            logger.error("数据库迁移失败")
            sys.exit(1)
    elif action == "rollback":
        success = rollback_migration()
        if success:
            logger.info("数据库回滚成功")
            sys.exit(0)
        else:
            logger.error("数据库回滚失败")
            sys.exit(1)
    else:
        print("无效的操作，请使用 'migrate' 或 'rollback'")
        sys.exit(1)


if __name__ == "__main__":
    main()